{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38fd0952",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "# 基于fenicsx框架 手动组装实现stokes方程\n",
    "方程请参考 `README.md`  \n",
    "速度组装使用fenicsx完成; 其余部分手动完成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "17e9239b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tools  # 实现的工具函数 用于有限元组装\n",
    "# import libraries\n",
    "from mpi4py import MPI\n",
    "import numpy as np\n",
    "\n",
    "import ufl\n",
    "from basix.ufl import element, mixed_element\n",
    "from dolfinx import default_real_type, fem, la\n",
    "from dolfinx.fem import (\n",
    "    Constant,\n",
    "    Function,\n",
    "    dirichletbc,\n",
    "    extract_function_spaces,\n",
    "    form,\n",
    "    functionspace,\n",
    "    locate_dofs_topological,\n",
    ")\n",
    "from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary\n",
    "from ufl import div, dx, grad, inner\n",
    "\n",
    "# PETSc\n",
    "# Demo文件给出了很多种情况的求解方案: 迭代求解 直接求解...\n",
    "try:\n",
    "    from petsc4py import PETSc\n",
    "\n",
    "    import dolfinx\n",
    "\n",
    "    if not dolfinx.has_petsc:\n",
    "        print(\"This demo requires DOLFINx to be compiled with PETSc enabled.\")\n",
    "        exit(0)\n",
    "    from petsc4py.PETSc import ScalarType\n",
    "except ModuleNotFoundError:\n",
    "    print(\"This demo requires petsc4py.\")\n",
    "    exit(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f26e0422",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义用到的函数\n",
    "import numpy as np\n",
    "def true_velocity(x):\n",
    "    u1 = x[0]**2 * x[1]**2 + np.exp(-x[1])\n",
    "    u2 = -2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])\n",
    "    return np.array([u1, u2])\n",
    "\n",
    "\n",
    "def true_pressure(x):\n",
    "    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1])\n",
    "\n",
    "\n",
    "def D(u):\n",
    "    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)\n",
    "\n",
    "\n",
    "def f1(x, y, nu):\n",
    "    \"\"\"\n",
    "    x方向载荷函数（NumPy版）\n",
    "    参数:\n",
    "        x: 标量或NumPy数组，x坐标\n",
    "        y: 标量或NumPy数组，y坐标\n",
    "        nu: 标量，粘性系数\n",
    "    返回:\n",
    "        标量或NumPy数组，计算结果\n",
    "    \"\"\"\n",
    "    r = -2 * nu * x**2 - 2 * nu * y**2 - nu * \\\n",
    "        np.exp(-y) + np.pi**2 * np.cos(np.pi * x) * np.cos(2 * np.pi * y)\n",
    "    return r\n",
    "\n",
    "\n",
    "def f2(x, y, nu):\n",
    "    \"\"\"\n",
    "    y方向载荷函数（NumPy版）\n",
    "    参数:\n",
    "        x: 标量或NumPy数组，x坐标\n",
    "        y: 标量或NumPy数组，y坐标\n",
    "        nu: 标量，粘性系数\n",
    "    返回:\n",
    "        标量或NumPy数组，计算结果\n",
    "    \"\"\"\n",
    "    r = 4 * nu * x * y \\\n",
    "        - nu * np.pi**3 * np.sin(np.pi * x) \\\n",
    "        + 2 * np.pi * (2 - np.pi * np.sin(np.pi * x)) * np.sin(2 * np.pi * y)\n",
    "    return r\n",
    "\n",
    "\n",
    "def getStiffnessMatrix(A: PETSc.Mat):\n",
    "    \"\"\"\n",
    "    获取刚度矩阵的值(np.array)\n",
    "    \"\"\"\n",
    "    m, n = A.getSize()\n",
    "    matrix_data = np.zeros((m, n))\n",
    "\n",
    "    # 获取矩阵的所有非零元素\n",
    "    row, col, val = A.getValuesCSR()\n",
    "\n",
    "    # 将CSR格式的数据转换为密集矩阵\n",
    "    for i in range(m):\n",
    "        for j in range(row[i], row[i + 1]):\n",
    "            matrix_data[i, col[j]] = val[j]\n",
    "\n",
    "    return matrix_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1e44a73e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[32, 32]剖分...\n"
     ]
    }
   ],
   "source": [
    "# 设置剖分\n",
    "nu = 1  # 粘度系数\n",
    "# idx = int(input(\"设置剖分(可选0, 1, 2, 3) = \"))  # 剖分索引\n",
    "idx = 1\n",
    "n_list = [[16, 16],\n",
    "          [32, 32],\n",
    "          [64, 64],\n",
    "          [128, 128],\n",
    "          [1, 1],\n",
    "          ]  # 剖分列表\n",
    "print(f\"{n_list[idx]}剖分...\")\n",
    "msh = create_rectangle(\n",
    "    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(\n",
    "        [1, 0])], n_list[idx], CellType.triangle\n",
    ")\n",
    "\n",
    "\n",
    "# Define Function Space: u1, u2 in V; p in Q\n",
    "P2 = element(\n",
    "    \"Lagrange\", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type\n",
    ")\n",
    "P1 = element(\"Lagrange\", msh.basix_cell(), degree=1, dtype=default_real_type)\n",
    "V = functionspace(msh, P2)\n",
    "Q = functionspace(msh, P1)\n",
    "\n",
    "\n",
    "facets = locate_entities_boundary(\n",
    "    msh, 1, lambda x: np.isclose(x[0], 1.0) | np.isclose(\n",
    "        x[0], 0.0) | np.isclose(x[1], -0.25) | np.isclose(x[1], 0.0)\n",
    ")\n",
    "dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)  # 边界点的自由度索引\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "60def14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 速度梯度内积 使用fenicsx实现\n",
    "# 定义变分问题 (关键修正：修正变分形式中的符号错误)\n",
    "(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)\n",
    "(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)\n",
    "\n",
    "a = form([[2 * nu * inner(D(u), D(v)) * dx, None],\n",
    "          [None, Constant(msh, PETSc.ScalarType(0)) * inner(p, q) * dx]])\n",
    "A = assemble_matrix_block(a)\n",
    "A.assemble()  # 组装速度梯度内积\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c0e6860",
   "metadata": {},
   "source": [
    "# 手动求解的核心代码  \n",
    "我们只组装B矩阵-(grad(u), q) [维持稳定性乘以-1]  \n",
    "B.T矩阵可以直接转置\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3f86c814",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_vector_basis_value(values):\n",
    "    # 解析向量基函数的值 尺寸为高斯点数 * 2 * (基函数 数 * 2)\n",
    "    res = np.zeros((values.shape[0], values.shape[2] // 2))\n",
    "    for i in range(values.shape[0]):\n",
    "        res[i] = values[i, 0, ::2]\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "bf7e1f45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 分配刚度矩阵和载荷向量尺寸\n",
    "num_cells = 2 * n_list[idx][0] * n_list[idx][1]  # 网格单元数量\n",
    "num_vel_dofs = 2 * (2*n_list[idx][0]+1)**2  # 速度自由度数量\n",
    "num_pressure_dofs = (n_list[idx][0]+1)**2  # 压力自由度数量\n",
    "num_basis_func_vel = P2.degree * 3  # 速度基函数数量\n",
    "num_basis_func_pressure = P1.degree * 3 # 压力基函数数量\n",
    "vel_dofmap = V.dofmap  # 速度自由度映射\n",
    "pressure_dofmap = Q.dofmap  # 压力自由度映射\n",
    "global_vertices = Q.tabulate_dof_coordinates()  # 网格所有节点坐标 按dof索引存储\n",
    "\n",
    "# 获取基函数值\n",
    "G = tools.GaussIntegrate(order=9)\n",
    "vel_basis_values_, vel_basis_values_d_xi, vel_basis_values_d_eta = P2.tabulate(\n",
    "    1, G.points)\n",
    "pressure_basis_values = P1.tabulate(0, G.points)[0]\n",
    "\n",
    "K = getStiffnessMatrix(A)  # 刚度矩阵\n",
    "F = np.zeros((num_vel_dofs + num_pressure_dofs,),\n",
    "             dtype=default_real_type)  # 载荷向量\n",
    "A = np.zeros((num_vel_dofs, num_vel_dofs))  # 速度矩阵\n",
    "B = np.zeros((num_pressure_dofs, num_vel_dofs),\n",
    "             dtype=default_real_type)  # B矩阵\n",
    "vel_basis_values = parse_vector_basis_value(vel_basis_values_)  # 解析向量基函数的值\n",
    "vel_basis_values_d_xi = parse_vector_basis_value(vel_basis_values_d_xi)\n",
    "vel_basis_values_d_eta = parse_vector_basis_value(vel_basis_values_d_eta)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3297fdb8",
   "metadata": {},
   "source": [
    "## 速度组装\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "    2K_1 + K_2 & K_3 \\\\  \n",
    "    K_4 & 2K_2 + K_1\n",
    "\\end{pmatrix}\n",
    "$$  \n",
    "这里显然也是对称的  我们利用$K_4$转置直接得到$K_3$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0cb7021b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 组装 -(grad(u), q) [维持稳定性最后整体乘以-1]\n",
    "for i in range(num_cells):\n",
    "    \n",
    "    vel_idx = vel_dofmap.cell_dofs(i)  # 速度自由度索引\n",
    "    pressure_idx = pressure_dofmap.cell_dofs(i)  # 压力自由度索引\n",
    "    # 本地单元的节点坐标\n",
    "    local_vertices = global_vertices[pressure_idx][:, :-1]\n",
    "    \n",
    "    bdt = tools.BasisDerivativeTransformer(local_vertices)\n",
    "    J_inv = np.linalg.inv(bdt.J.T)\n",
    "    det_J = bdt.detJ\n",
    "\n",
    "    gs_points_phy = G.to_physical_gauss_points(*local_vertices)  # 物理高斯点\n",
    "    \n",
    "    \n",
    "    # 组装速度内积项\n",
    "    for alpha in range(num_basis_func_vel):\n",
    "        for beta in range(num_basis_func_vel):\n",
    "            local_diff_xi = bdt.local_dx1(vel_basis_values_d_xi, vel_basis_values_d_eta)\n",
    "            local_diff_eta = bdt.local_dy1(vel_basis_values_d_xi, vel_basis_values_d_eta)\n",
    "            item_1 = local_diff_xi[:, alpha] * local_diff_xi[:, beta]\n",
    "            item_2 = local_diff_eta[:, alpha] * local_diff_eta[:, beta]\n",
    "            item_4 = local_diff_xi[:, alpha] * local_diff_eta[:, beta]\n",
    "            K1 = nu * det_J * np.dot(item_1, G.weights)\n",
    "            K2 = nu * det_J * np.dot(item_2, G.weights)\n",
    "            K4 = nu * det_J * np.dot(item_4, G.weights)\n",
    "\n",
    "            # 组装\n",
    "            A[2 * vel_idx[beta], 2 * vel_idx[alpha]] += 2*K1 + K2\n",
    "            A[2 * vel_idx[beta]+1, 2 * vel_idx[alpha]+1] += K1 + 2*K2\n",
    "            A[2 * vel_idx[beta], 2 * vel_idx[alpha]+1] += K4\n",
    "            A[2 * vel_idx[beta]+1, 2 * vel_idx[alpha]] += K4\n",
    "            \n",
    "    for alpha in range(num_basis_func_vel):\n",
    "        for beta in range(num_basis_func_pressure):\n",
    "            # local基函数的导数 存储x和y方向的导数 分别记录索引为alpha, beta的基函数\n",
    "            local_diff_alpha = J_inv @ np.vstack(\n",
    "                (vel_basis_values_d_xi[:, alpha], vel_basis_values_d_eta[:, alpha]))            \n",
    "            \n",
    "            item_x = -local_diff_alpha[0] * pressure_basis_values[:, beta]\n",
    "            item_y = -local_diff_alpha[1] * pressure_basis_values[:, beta]\n",
    "            B[pressure_idx[beta], 2 * vel_idx[alpha]] += np.dot(item_x, G.weights) * det_J\n",
    "            B[pressure_idx[beta], 2 * vel_idx[alpha] +1] += np.dot(item_y, G.weights) * det_J\n",
    "             \n",
    "            \n",
    "    # 组装f 载荷向量\n",
    "    for beta in range(num_basis_func_vel):\n",
    "        item_x = f1(gs_points_phy[:, 0], gs_points_phy[:, 1], nu=nu) * vel_basis_values[:, beta]\n",
    "        item_y = f2(gs_points_phy[:, 0], gs_points_phy[:, 1], nu=nu) * vel_basis_values[:, beta]\n",
    "        F[2 * vel_idx[beta]] += np.dot(item_x, G.weights) * det_J\n",
    "        F[2 * vel_idx[beta] + 1] += np.dot(item_y, G.weights) * det_J\n",
    "\n",
    "# 将计算结果写入刚度矩阵\n",
    "K = tools.assign_to_corner(K, A, corner=\"upper_left\")\n",
    "K = tools.assign_to_corner(K, B, corner=\"lower_left\")\n",
    "K = tools.assign_to_corner(K, B.T, corner=\"upper_right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "f3624f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置边界条件\n",
    "v_dofs = fem.locate_dofs_topological(\n",
    "    V=V, entity_dim=1, entities=facets)  # 速度的边界节点\n",
    "p_dofs = fem.locate_dofs_topological(\n",
    "    V=Q, entity_dim=1, entities=facets)  # 压力的边界节点\n",
    "for v_dof in v_dofs:\n",
    "    boundary_coords = V.tabulate_dof_coordinates()[v_dof][:-1]\n",
    "    K[2 * v_dof] = 0\n",
    "    K[2 * v_dof + 1] = 0\n",
    "    F[2 * v_dof], F[2 * v_dof + 1] = true_velocity(boundary_coords)  # 速度边界值\n",
    "    \n",
    "    K[2 * v_dof, 2 * v_dof] = 1\n",
    "    K[2 * v_dof + 1, 2 * v_dof + 1] = 1\n",
    "# 固定压力为0\n",
    "# idx==13 对应 (0, 0)坐标 固定压力 便于测试\n",
    "fixed_idx = 0\n",
    "p_fix_x, p_fix_y = Q.tabulate_dof_coordinates()[fixed_idx][:-1]  # 固定压力点的坐标\n",
    "K[num_vel_dofs + fixed_idx] = 0\n",
    "K[num_vel_dofs + fixed_idx, num_vel_dofs + fixed_idx] = 1\n",
    "\n",
    "F[num_vel_dofs + fixed_idx] = true_pressure([p_fix_x, p_fix_y])\n",
    "\n",
    "# solve\n",
    "res = np.linalg.solve(K, F)  # 求解线性方程组\n",
    "\n",
    "np.save(\"K.npy\", K)  # 保存刚度矩阵\n",
    "np.save(\"F.npy\", F)  # 保存载荷向量\n",
    "\n",
    "u_sol, p_sol = res[:num_vel_dofs], res[num_vel_dofs:]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "759a4451",
   "metadata": {},
   "source": [
    "# 误差计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2820f20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def u1_dx(x, y):\n",
    "    return 2 * x * y**2\n",
    "    \n",
    "def u1_dy(x, y):\n",
    "    return 2 * x**2 * y - np.exp(-y)\n",
    "    \n",
    "def u2_dx(x, y):\n",
    "    return -2/3 * y**3 - np.pi**2 * np.cos(np.pi * x)\n",
    "    \n",
    "def u2_dy(x, y):\n",
    "    return -2 * x * y**2 \n",
    "    \n",
    "def p_dx(x, y):\n",
    "    return np.pi**2 * np.cos(np.pi * x) * np.cos(2 * np.pi * y)\n",
    "    \n",
    "def p_dy(x, y):\n",
    "    return(4 * np.pi - 2 * np.pi**2 * np.sin(np.pi * x)) * np.sin(2 * np.pi * y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "08887b4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================n=32误差===============\n",
      "Max velocity err: 4.50109703065138e-07\n",
      "Max pressure err: 0.0025918318412825325\n",
      "\n",
      "L2 velocity err: 5.812499500204085e-06\n",
      "L2 pressure err: 0.0007432798117151815\n",
      "\n",
      "H1 velocity err: 0.001253428567368367\n",
      "H1 pressure err: 0.21777522649151587\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import norm\n",
    "\n",
    "u_L2_err = norm.get_L2_err(\n",
    "    msh=msh,\n",
    "    dof_map=vel_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func=true_velocity,\n",
    "    u_num=u_sol,\n",
    "    func_space=V,\n",
    "    order=2,\n",
    ")\n",
    "p_l2_err = norm.get_L2_err(\n",
    "    msh=msh,\n",
    "    dof_map=pressure_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func=true_pressure,\n",
    "    u_num=p_sol,\n",
    "    func_space=Q,\n",
    "    order=1,\n",
    ")\n",
    "u_max_err = norm.get_max_node_err(\n",
    "    u_true_func=true_velocity,\n",
    "    u_num=u_sol,\n",
    "    func_space=V,\n",
    ")\n",
    "p_max_err = norm.get_max_node_err(\n",
    "    u_true_func=true_pressure,\n",
    "    u_num=p_sol,\n",
    "    func_space=Q,\n",
    ")\n",
    "u_h1_err = norm.get_H1_err(\n",
    "    msh=msh,\n",
    "    dof_map=vel_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func_diff_ls=[u1_dx, u1_dy, u2_dx, u2_dy],\n",
    "    u_num=u_sol,\n",
    "    func_space=V,\n",
    "    order=2,\n",
    ")\n",
    "p_h1_err = norm.get_H1_err(\n",
    "    msh=msh,\n",
    "    dof_map=pressure_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func_diff_ls=[p_dx, p_dy],\n",
    "    u_num=p_sol,\n",
    "    func_space=Q,\n",
    "    order=1,\n",
    ")\n",
    "\n",
    "\n",
    "print(f\"================n={n_list[idx][0]}误差===============\")\n",
    "print(f\"Max velocity err: {u_max_err}\")\n",
    "print(f\"Max pressure err: {p_max_err}\\n\")\n",
    "\n",
    "print(f\"L2 velocity err: {u_L2_err}\")\n",
    "print(f\"L2 pressure err: {p_l2_err}\\n\")\n",
    "\n",
    "print(f\"H1 velocity err: {u_h1_err}\")\n",
    "print(f\"H1 pressure err: {p_h1_err}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eee3597",
   "metadata": {},
   "source": [
    "================n=16误差===============  \n",
    "Max velocity err: 6.868874497634181e-06  \n",
    "Max pressure err: 0.010482784426438203  \n",
    "\n",
    "L2 velocity err: 4.644321399673993e-05  \n",
    "L2 pressure err: 0.002974667749951565  \n",
    "\n",
    "H1 velocity err: 0.0050125017372174  \n",
    "H1 pressure err: 0.4360833760639439  \n",
    "\n",
    "==================n = 32误差=================  \n",
    "Max velocity err: 4.50109703065138e-07  \n",
    "Max pressure err: 0.0025918318412825325  \n",
    "\n",
    "L2 velocity err: 5.812499500204085e-06  \n",
    "L2 pressure err: 0.0007432798117151815  \n",
    "\n",
    "H1 velocity err: 0.001253428567368367  \n",
    "H1 pressure err: 0.21777522649151587"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "21d0f5a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# deprecated\n",
    "# 请使用norms.py中的getL2_err函数\n",
    "def getL2_err(dof_map, num_cells, u_true_func, u_num, basis_values, func_space, type_=\"scalar\"):\n",
    "    \"\"\"给定单元 返回单元的L2误差\n",
    "\n",
    "    Args:\n",
    "        dof_map (_type_): 自由度索引表\n",
    "        num_cells (_type_): 单元id总数\n",
    "        u_true_func (_type_): 真解的函数\n",
    "        u_num (_type_): 变量数值解\n",
    "        basis_values: _type_: 标准单元的基函数值\n",
    "        func_space: 基函数空间\n",
    "        type_ (str, optional): 处理scalar或vector. Defaults to \"scalar\".\n",
    "    \"\"\"\n",
    "    L2_err = 0.0  # L2误差平方和\n",
    "    global_dofs = func_space.tabulate_dof_coordinates()  # 网格所有节点坐标 按dof索引存储\n",
    "    num_gs_points = basis_values.shape[0]  # 高斯点数量\n",
    "    num_basis_func = basis_values.shape[1]  # 基函数数量\n",
    "    G = tools.GaussIntegrate(order=num_gs_points)  # 高斯积分对象\n",
    "    for cell_idx in range(num_cells):\n",
    "    # 计算基函数的值\n",
    "        dof_idx = dof_map.cell_dofs(cell_idx)  # 获取自由度索引\n",
    "        local_nodes = global_dofs[dof_idx][:, :-1]  # 自由度节点的全局坐标\n",
    "        gs_phy_points = G.to_physical_gauss_points(*local_nodes[:3])\n",
    "        bdt = tools.BasisDerivativeTransformer(local_vertices)\n",
    "        det_J = bdt.detJ  # 雅可比行列式\n",
    "        # 数值解\n",
    "        if type_ == \"scalar\":\n",
    "            u_nodes_num = u_num[dof_idx]  # 数值解的节点值\n",
    "            u_gs_true = u_true_func(gs_phy_points.T)  # 真解在高斯点的值\n",
    "            u_gs_num = np.zeros_like(u_gs_true)\n",
    "            for alpha in range(num_gs_points):\n",
    "                # 计算高斯点的物理坐标\n",
    "                u_gs_num[alpha] = np.dot(basis_values[alpha], u_nodes_num)\n",
    "            r1 = (u_gs_num - u_gs_true) ** 2  # 误差平方\n",
    "            L2_err += np.dot(r1, G.weights) * det_J\n",
    "        if type_ == \"vector\":\n",
    "            # 速度分量的计算\n",
    "            u_x_nodes_num = u_num[2 * dof_idx]\n",
    "            u_y_nodes_num = u_num[2 * dof_idx + 1]\n",
    "            u_x_gs_true, u_y_gs_true = u_true_func(gs_phy_points.T)  # 真解\n",
    "            u_x_gs_num, u_y_gs_num = np.zeros_like(u_x_gs_true), np.zeros_like(u_y_gs_true)\n",
    "            for alpha in range(num_gs_points):\n",
    "                u_x_gs_num[alpha] = np.dot(basis_values[alpha], u_x_nodes_num)\n",
    "                u_y_gs_num[alpha] = np.dot(basis_values[alpha], u_y_nodes_num)\n",
    "            r1 = (u_x_gs_num - u_x_gs_true) ** 2 + (u_y_gs_num - u_y_gs_true) ** 2\n",
    "            L2_err += np.dot(r1, G.weights) * det_J\n",
    "    return np.sqrt(L2_err)  # 返回L2误差平方根\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "fcfc5a97",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pressure L2 error: 0.0007432798117151815\n",
      "Velocity L2 error: 5.812499500204085e-06\n"
     ]
    }
   ],
   "source": [
    "pressure_L2 = getL2_err(\n",
    "    dof_map=pressure_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func=true_pressure,\n",
    "    u_num=p_sol,\n",
    "    basis_values=pressure_basis_values,\n",
    "    func_space=Q,\n",
    "    type_=\"scalar\"\n",
    ")\n",
    "vel_L2 = getL2_err(\n",
    "    dof_map=vel_dofmap,\n",
    "    num_cells=num_cells,\n",
    "    u_true_func=true_velocity,\n",
    "    u_num=u_sol,\n",
    "    basis_values=vel_basis_values,\n",
    "    func_space=V,\n",
    "    type_=\"vector\"\n",
    ")\n",
    "print(f\"Pressure L2 error: {pressure_L2}\")\n",
    "print(f\"Velocity L2 error: {vel_L2}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
